from datetime import datetime, timedelta

from config import Config
from extensions import db
from logger_config import get_logger
from models import RequestRecord, CostRecord, ModelPricing
from utils.token_utils import estimate_token_count

logger = get_logger(__name__)


class CostService:
    """成本服务类，负责计算和管理查询成本"""

    def __init__(self):
        self.config = Config()
        # 默认定价（每1000个token的价格，单位：美元）
        self.default_pricing = {
            'claude-3-5-sonnet-20241022': {
                'input': 0.003,  # $0.003 per 1K input tokens
                'output': 0.015  # $0.015 per 1K output tokens
            },
            'claude-3-5-haiku-20241022': {
                'input': 0.001,  # $0.001 per 1K input tokens
                'output': 0.005  # $0.005 per 1K output tokens
            },
            'gpt-4o': {
                'input': 0.005,  # $0.005 per 1K input tokens
                'output': 0.015  # $0.015 per 1K output tokens
            },
            'gpt-4o-mini': {
                'input': 0.00015,  # $0.00015 per 1K input tokens
                'output': 0.0006  # $0.0006 per 1K output tokens
            },
            'gpt-3.5-turbo': {
                'input': 0.001,  # $0.001 per 1K input tokens
                'output': 0.002  # $0.002 per 1K output tokens
            },
            'simple_response': {
                'input': 0.0001,  # $0.0001 per 1K input tokens (低成本简单对话)
                'output': 0.0002  # $0.0002 per 1K output tokens (低成本简单对话)
            }
        }

    def calculate_query_cost(self, query_text, response_text, execution_summary=None, model_name=None):
        """计算查询成本
        
        Args:
            query_text: 查询文本
            response_text: 响应文本
            execution_summary: 执行摘要（可选）
            model_name: 模型名称（可选）
        
        Returns:
            dict: 成本信息
        """
        try:
            # 获取模型名称
            if not model_name:
                model_name = self.config.DEFAULT_MODEL

            # 估算token数量
            input_tokens = self._estimate_input_tokens(query_text)
            output_tokens = self._estimate_output_tokens(response_text, execution_summary)

            # 获取定价信息
            pricing = self._get_model_pricing(model_name)

            # 计算成本
            input_cost = (input_tokens / 1000) * pricing['input']
            output_cost = (output_tokens / 1000) * pricing['output']
            total_cost = input_cost + output_cost

            cost_info = {
                'model_name': model_name,
                'input_tokens': input_tokens,
                'output_tokens': output_tokens,
                'total_tokens': input_tokens + output_tokens,
                'input_cost': round(input_cost, 6),
                'output_cost': round(output_cost, 6),
                'total_cost': round(total_cost, 6),
                'pricing': pricing,
                'currency': 'USD'
            }

            logger.info(f"成本计算完成 - 模型: {model_name}, 总成本: ${total_cost:.6f}")

            return cost_info

        except Exception as e:
            logger.error(f"计算查询成本失败: {str(e)}")
            return {
                'model_name': model_name or 'unknown',
                'input_tokens': 0,
                'output_tokens': 0,
                'total_tokens': 0,
                'input_cost': 0.0,
                'output_cost': 0.0,
                'total_cost': 0.0,
                'pricing': {'input': 0.0, 'output': 0.0},
                'currency': 'USD',
                'error': str(e)
            }

    def _estimate_input_tokens(self, query_text):
        """估算输入token数量
        
        Args:
            query_text: 查询文本
        
        Returns:
            int: 估算的token数量
        """
        try:
            return estimate_token_count(query_text)
        except Exception as e:
            logger.warning(f"估算输入token失败，使用默认方法: {str(e)}")
            # 简单估算：平均每个单词约1.3个token
            return max(1, int(len(query_text.split()) * 1.3))

    def _estimate_output_tokens(self, response_text, execution_summary=None):
        """估算输出token数量
        
        Args:
            response_text: 响应文本
            execution_summary: 执行摘要（可选）
        
        Returns:
            int: 估算的token数量
        """
        try:
            total_tokens = 0

            # 估算响应文本的token
            if response_text:
                total_tokens += estimate_token_count(str(response_text))

            # 如果有执行摘要，也计算其token
            if execution_summary and isinstance(execution_summary, dict):
                summary_text = str(execution_summary.get('summary', ''))
                if summary_text:
                    total_tokens += estimate_token_count(summary_text)

            return max(1, total_tokens)

        except Exception as e:
            logger.warning(f"估算输出token失败，使用默认方法: {str(e)}")
            # 简单估算：平均每个单词约1.3个token
            text_length = len(str(response_text).split()) if response_text else 0
            return max(1, int(text_length * 1.3))

    def _get_model_pricing(self, model_name):
        """获取模型定价信息
        
        Args:
            model_name: 模型名称
        
        Returns:
            dict: 定价信息
        """
        try:
            # 首先尝试从数据库中获取
            pricing = ModelPricing.query.filter_by(model_name=model_name, is_active=True).first()
            if pricing:
                return {
                    'input': pricing.input_price_per_1k,
                    'output': pricing.output_price_per_1k
                }
            
            # 然后尝试从配置中获取
            if hasattr(self.config, 'MODEL_PRICING') and model_name in self.config.MODEL_PRICING:
                return self.config.MODEL_PRICING[model_name]

            # 最后尝试从默认定价中获取（作为后备）
            if model_name in self.default_pricing:
                return self.default_pricing[model_name]

            # 如果找不到，使用默认定价
            logger.warning(f"未找到模型 {model_name} 的定价信息，使用默认定价")
            return {
                'input': 0.001,  # 默认输入价格
                'output': 0.002  # 默认输出价格
            }
        except Exception as e:
            logger.error(f"获取模型定价失败: {str(e)}，使用默认定价")
            return {
                'input': 0.001,  # 默认输入价格
                'output': 0.002  # 默认输出价格
            }

    def get_user_cost_summary(self, user_id, days=30):
        """获取用户成本摘要
        
        Args:
            user_id: 用户ID
            days: 统计天数
        
        Returns:
            dict: 成本摘要
        """
        try:
            # 计算时间范围
            end_date = datetime.utcnow()
            start_date = end_date - timedelta(days=days)

            # 查询用户在指定时间范围内的查询记录
            queries = db.session.query(RequestRecord).filter(
                RequestRecord.user_id == user_id,
                RequestRecord.created_at >= start_date,
                RequestRecord.created_at <= end_date,
                RequestRecord.status == 'completed'
            ).all()

            # 统计成本信息
            total_cost = 0.0
            total_queries = len(queries)
            total_input_tokens = sum(q.input_tokens or 0 for q in queries)
            total_output_tokens = sum(q.output_tokens or 0 for q in queries)
            total_tokens = total_input_tokens + total_output_tokens
            
            # 计算总成本（通过关联的CostRecord）
            for query in queries:
                if query.cost_record:
                    total_cost += query.cost_record.cost

            # 按模型统计
            model_stats = {}
            for query in queries:
                # 从查询记录中提取模型信息
                model = query.model if query.model else 'unknown'

                if model not in model_stats:
                    model_stats[model] = {
                        'queries': 0,
                        'cost': 0.0,
                        'input_tokens': 0,
                        'output_tokens': 0
                    }

                model_stats[model]['queries'] += 1
                if query.cost_record:
                    model_stats[model]['cost'] += query.cost_record.cost
                model_stats[model]['input_tokens'] += query.input_tokens or 0
                model_stats[model]['output_tokens'] += query.output_tokens or 0

            # 计算平均成本
            avg_cost_per_query = total_cost / total_queries if total_queries > 0 else 0
            avg_tokens_per_query = total_tokens / total_queries if total_queries > 0 else 0

            return {
                'period_days': days,
                'start_date': start_date.isoformat(),
                'end_date': end_date.isoformat(),
                'total_cost': round(total_cost, 6),
                'total_queries': total_queries,
                'total_input_tokens': total_input_tokens,
                'total_output_tokens': total_output_tokens,
                'total_tokens': total_tokens,
                'avg_cost_per_query': round(avg_cost_per_query, 6),
                'avg_tokens_per_query': round(avg_tokens_per_query, 2),
                'model_stats': model_stats,
                'currency': 'USD'
            }

        except Exception as e:
            logger.error(f"获取用户成本摘要失败: {str(e)}")
            return {
                'period_days': days,
                'total_cost': 0.0,
                'total_queries': 0,
                'total_input_tokens': 0,
                'total_output_tokens': 0,
                'total_tokens': 0,
                'avg_cost_per_query': 0.0,
                'avg_tokens_per_query': 0.0,
                'model_stats': {},
                'currency': 'USD',
                'error': str(e)
            }

    def get_system_cost_summary(self, days=30):
        """获取系统成本摘要
        
        Args:
            days: 统计天数
        
        Returns:
            dict: 系统成本摘要
        """
        try:
            # 计算时间范围
            end_date = datetime.utcnow()
            start_date = end_date - timedelta(days=days)

            # 查询所有用户在指定时间范围内的查询记录
            queries = db.session.query(RequestRecord).filter(
                RequestRecord.created_at >= start_date,
                RequestRecord.created_at <= end_date,
                RequestRecord.status == 'completed'
            ).all()

            # 统计总体信息
            total_cost = 0.0
            total_queries = len(queries)
            total_users = len(set(q.user_id for q in queries))
            total_input_tokens = sum(q.input_tokens or 0 for q in queries)
            total_output_tokens = sum(q.output_tokens or 0 for q in queries)
            total_tokens = total_input_tokens + total_output_tokens
            
            # 计算总成本（通过关联的CostRecord）
            for query in queries:
                if query.cost_record:
                    total_cost += query.cost_record.cost

            # 按用户统计
            user_stats = {}
            for query in queries:
                user_id = query.user_id
                if user_id not in user_stats:
                    user_stats[user_id] = {
                        'queries': 0,
                        'cost': 0.0,
                        'tokens': 0
                    }

                user_stats[user_id]['queries'] += 1
                if query.cost_record:
                    user_stats[user_id]['cost'] += query.cost_record.cost
                user_stats[user_id]['tokens'] += (query.input_tokens or 0) + (query.output_tokens or 0)

            # 找出成本最高的用户
            top_users = sorted(
                user_stats.items(),
                key=lambda x: x[1]['cost'],
                reverse=True
            )[:10]  # 前10名

            return {
                'period_days': days,
                'start_date': start_date.isoformat(),
                'end_date': end_date.isoformat(),
                'total_cost': round(total_cost, 6),
                'total_queries': total_queries,
                'total_users': total_users,
                'total_input_tokens': total_input_tokens,
                'total_output_tokens': total_output_tokens,
                'total_tokens': total_tokens,
                'avg_cost_per_query': round(total_cost / total_queries, 6) if total_queries > 0 else 0,
                'avg_cost_per_user': round(total_cost / total_users, 6) if total_users > 0 else 0,
                'avg_queries_per_user': round(total_queries / total_users, 2) if total_users > 0 else 0,
                'top_users_by_cost': [{
                    'user_id': user_id,
                    'queries': stats['queries'],
                    'cost': round(stats['cost'], 6),
                    'tokens': stats['tokens']
                } for user_id, stats in top_users],
                'currency': 'USD'
            }

        except Exception as e:
            logger.error(f"获取系统成本摘要失败: {str(e)}")
            return {
                'period_days': days,
                'total_cost': 0.0,
                'total_queries': 0,
                'total_users': 0,
                'total_input_tokens': 0,
                'total_output_tokens': 0,
                'total_tokens': 0,
                'avg_cost_per_query': 0.0,
                'avg_cost_per_user': 0.0,
                'avg_queries_per_user': 0.0,
                'top_users_by_cost': [],
                'currency': 'USD',
                'error': str(e)
            }

    def update_model_pricing(self, model_name, input_price, output_price):
        """更新模型定价
        
        Args:
            model_name: 模型名称
            input_price: 输入价格（每1000个token）
            output_price: 输出价格（每1000个token）
        
        Returns:
            dict: 更新结果
        """
        try:
            # 验证价格
            if input_price < 0 or output_price < 0:
                return {
                    'success': False,
                    'error': '价格不能为负数',
                    'error_code': 'INVALID_PRICE'
                }

            # 更新默认定价
            self.default_pricing[model_name] = {
                'input': input_price,
                'output': output_price
            }

            logger.info(f"模型定价更新成功: {model_name} - 输入: ${input_price}, 输出: ${output_price}")

            return {
                'success': True,
                'message': f'模型 {model_name} 定价更新成功'
            }

        except Exception as e:
            logger.error(f"更新模型定价失败: {str(e)}")
            return {
                'success': False,
                'error': f'更新定价失败: {str(e)}',
                'error_code': 'PRICING_UPDATE_ERROR'
            }

    def get_all_model_pricing(self):
        """获取所有模型定价
        
        Returns:
            dict: 所有模型的定价信息
        """
        try:
            # 合并配置中的定价和默认定价
            all_pricing = self.default_pricing.copy()

            if hasattr(self.config, 'MODEL_PRICING'):
                all_pricing.update(self.config.MODEL_PRICING)

            return {
                'success': True,
                'pricing': all_pricing,
                'currency': 'USD'
            }

        except Exception as e:
            logger.error(f"获取模型定价失败: {str(e)}")
            return {
                'success': False,
                'pricing': {},
                'currency': 'USD',
                'error': str(e)
            }


# 全局成本服务实例
cost_service = CostService()
